import torch
import torch.nn as nn

import math
from .module import LAE_SE, SALayer, CALayer, DWT, IWT
import torch.nn.functional as F


class FusionModule(nn.Module):
    def __init__(self, channel, numclass, sptial, model_num=3):
        super(FusionModule, self).__init__()
        self.total_feature_maps = {}

        self.conv1 = nn.Conv2d(channel * model_num, channel * model_num, kernel_size=3, stride=1, padding=1, groups=channel*model_num, bias=False)
        self.bn1 = nn.BatchNorm2d(channel * model_num)
        self.relu1 = nn.ReLU(True)
        self.conv1_1 = nn.Conv2d(channel * model_num, channel, kernel_size=1, groups=1, bias=False)
        self.bn1_1 = nn.BatchNorm2d(channel)
        self.relu1_1 = nn.ReLU(True)

        self.avgpool2d = nn.AvgPool2d(sptial)
        self.lae = LAE_SE(channel)
        self.fc2 = nn.Linear(channel, numclass)


        self.sptial = sptial

        self.dim_adjust = nn.Sequential(  # 使用Sequential容器构建更复杂的降维层
            nn.Linear(40960, 8192),  # 加入中间层，逐步降维
            nn.BatchNorm1d(8192),  # 加入批归一化
            nn.ReLU(True),  # 加入ReLU激活函数
            nn.Linear(8192, 2048),  # 加入中间层，逐步降维
            nn.BatchNorm1d(2048),  # 加入批归一化
            nn.ReLU(True),  # 加入ReLU激活函数
            nn.Linear(2048, 512),  # 进一步降维到512
            nn.BatchNorm1d(512),  # 加入批归一化
            nn.ReLU(True)  # 加入ReLU激活函数
        )
        # for m in self.modules():
        #     print(type(m))

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
        self.register_hook()

    def forward(self, x):

        x = self.relu1(self.bn1((self.conv1(x))))
        x = self.relu1_1(self.bn1_1(self.conv1_1(x)))

        # x = self.avgpool2d(x)
        lae_out = self.lae(x)
        avg_out = self.avgpool2d(x)
        # print(lae_out)
        # print(avg_out)
        # ds = nn.AdaptiveAvgPool2d((1,1))
        upsample_bilinear = nn.Upsample(size=(4, 4), mode='bilinear', align_corners=False)
        upsampled_tensor2_bilinear = upsample_bilinear(avg_out)
        x = torch.cat([lae_out, upsampled_tensor2_bilinear], dim=1)
        x = x.view(x.size(0), -1)
        x = self.dim_adjust(x)
        out = self.fc2(x)

        return out

    def register_hook(self):

        self.extract_layers = ['relu1_1']

        def get_activation(maps, name):
            def get_output_hook(module, input, output):
                maps[name+str(output.device)] = output

            return get_output_hook

        def add_hook(model, maps, extract_layers):
            for name, module in model.named_modules():
                if name in extract_layers:
                    module.register_forward_hook(get_activation(maps, name))

        add_hook(self, self.total_feature_maps, self.extract_layers)


class TransFusionModule(nn.Module):

    def __init__(self, channel, sptial, model_num=3):
        super(TransFusionModule, self).__init__()

        self.conv1 = nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1, groups=channel, bias=False)
        self.bn1 = nn.BatchNorm2d(channel)
        self.relu1 = nn.ReLU(True)
        self.conv1_1 = nn.Conv2d(channel, channel * model_num, kernel_size=1, groups=1, bias=False)
        self.bn1_1 = nn.BatchNorm2d(channel * model_num)
        self.relu1_1 = nn.ReLU(True)
        self.sa = SALayer()
        self.ca = CALayer(channel)
        self.dwt = DWT()  # 添加小波变换层
        self.iwt = IWT()  # 添加逆小波变换层
        self.conv_fusion = nn.Conv2d(channel * model_num * 2, channel * model_num, kernel_size=1, bias=False)  # 融合卷积


        self.sptial = sptial

    def forward(self, input):
        device = input.device
        x = self.relu1(self.bn1((self.conv1(input))))
        x = self.sa(x)
        x = self.ca(x)
        # 特征融合部分
        x_dwt = self.dwt(x.to(device))  # 小波变换
        x_iwt = self.iwt(x_dwt.to(device))  # 逆小波变换(简化版，实际可能需要根据你的DWT调整)
        # 这里假设你的DWT输出是channel*2
        original_out = self.relu1_1(self.bn1_1(self.conv1_1(x)))


        if original_out.device != device:
            original_out = original_out.to(device)
        if x_iwt.device != device:
            x_iwt = x_iwt.to(device)
        x_iwt = F.interpolate(x_iwt, size=(7, 7), mode='bicubic', align_corners=False)
        fused_features = torch.cat([original_out, x_iwt], dim=1)  # 拼接
        # 计算conv_fusion的输入通道数
        num_input_channels = fused_features.shape[1]
        # 动态创建conv_fusion层，根据实际输入通道数调整
        self.conv_fusion = nn.Conv2d(num_input_channels, self.conv1_1.out_channels, kernel_size=1, bias=False).to(device)
        out = self.conv_fusion(fused_features)
        out = self.relu1_1(self.bn1_1(out))  # 融合卷积

        return out.to(device)
